"""
写这个组件是为了什么？

1、加载我们的Excel数据，到一个统一的类里面；
2、方便后续能快速的获取到数据源的结构信息（get_data_info）
3、方便直接执行SQL代码出结果；

database是用来创建一个小型的SQL数据库，便于后续执行SQL代码

"""

from typing import Optional, Union, List, Dict
import pandas as pd
import sqlite3
import os
import re
from typing import Optional, List
from nl2sql.datasource.base import BaseDatasource, Table, Column
from nl2sql.model.embedding_model import EmbeddingModel


class ExcelDatasource(BaseDatasource):
    """Excel文件的数据源类，继承自BaseDatasource，支持自动创建SQLite数据库"""

    def __init__(
            self,
            source_path: str,
            auto_create_sqlite: bool = True,
            db_path: Optional[str] = None,
            embedding_model: EmbeddingModel = None,
    ):
        """
        初始化Excel数据源，支持单个Excel文件或包含多个Excel的文件夹

        Args:
            source_path (str): Excel文件路径或目录路径
            auto_create_sqlite (bool): 是否自动创建SQLite数据库
            db_path (Optional[str]): SQLite数据库路径，默认为内存数据库
        """
        super().__init__(source_path=source_path, embedding_model=embedding_model)
        self.auto_create_sqlite = auto_create_sqlite
        self.db_path = db_path if db_path else ":memory:"
        self.connection = sqlite3.connect(self.db_path)
        self.all_table_info = {}
        self.load_data()
        self.initialize_enums()

    def _clean_table_name(self, name: str) -> str:
        """清理表名，替换非法字符，确保符合SQLite命名规范"""
        return re.sub(r'[\s-]', '_', name).strip()

    def _create_table_metadata(self, table_name: str, df: pd.DataFrame) -> Table:
        """创建表的元数据，包含列信息"""
        columns = [
            Column(
                name_en=col,
                name_cn=col,  # 暂时使用英文名，后续可扩展为中文映射
                enums=df[col].dropna().unique().tolist() if df[col].dtype in ['object', 'category'] else []
            )
            for col in df.columns
        ]
        return Table(
            name_en=table_name,
            name_cn=table_name,  # 暂时使用英文名，后续可扩展为中文映射
            columns=columns
        )

    def _process_excel_file(self, file_path: str, file_base: str) -> None:
        """处理单个Excel文件，将每个sheet转换为SQLite表"""
        try:
            sheets = pd.read_excel(file_path, sheet_name=None)
            for sheet_name, df in sheets.items():
                table_name = self._clean_table_name(f"{file_base}_{sheet_name}")
                df.to_sql(table_name, self.connection, if_exists="replace", index=False)
                self.all_table_info[table_name] = list(df.columns)
                self.tables.append(self._create_table_metadata(table_name, df))
                print(f"✅ 加载表: {table_name}")
        except Exception as e:
            print(f"❌ 加载失败: {file_path}, 错误: {e}")

    def load_data(self) -> None:
        """
        加载Excel文件或目录下所有Excel文件，将每个sheet转换为SQLite表，并更新tables元数据。
        """
        if not os.path.exists(self.source_path):
            raise FileNotFoundError(f"路径不存在: {self.source_path}")

        if os.path.isdir(self.source_path):
            # 目录模式
            excel_files = [
                f for f in os.listdir(self.source_path)
                if f.endswith((".xlsx", ".xls"))
            ]
            if not excel_files:
                print(f"⚠️ 目录 {self.source_path} 中未找到Excel文件")
                return

            for file_name in excel_files:
                file_path = os.path.join(self.source_path, file_name)
                file_base = os.path.splitext(file_name)[0]
                self._process_excel_file(file_path, file_base)
        else:
            # 单文件模式
            if not self.source_path.endswith((".xlsx", ".xls")):
                raise ValueError(f"文件格式不支持: {self.source_path}")
            file_base = os.path.splitext(os.path.basename(self.source_path))[0]
            self._process_excel_file(self.source_path, file_base)
            self.data = pd.read_excel(self.source_path)  # 仅为单文件模式设置self.data

    def get_available_sheets(self) -> List[str]:
        """
        获取Excel文件中所有sheet的名称

        Returns:
            List[str]: sheet名称列表
        """
        with pd.ExcelFile(self.source_path) as xls:
            return xls.sheet_names

    def get_data_info(self) -> str:

        info_text = "以下是数据源中的表结构和信息：\n\n"

        for table_name in self.all_table_info.keys():
            # 使用 pandas 获取数据类型和样本数据
            df = pd.read_sql_query(f"SELECT * FROM {table_name} LIMIT 3", self.connection)

            # 构造列描述
            columns_desc = []
            for col_name in df.columns:
                dtype = str(df[col_name].dtype)
                inferred_type = "TEXT"  # 默认类型
                if "int" in dtype:
                    inferred_type = "INTEGER"
                elif "float" in dtype:
                    inferred_type = "REAL"
                elif "bool" in dtype:
                    inferred_type = "BOOLEAN"
                elif "datetime" in dtype:
                    inferred_type = "DATETIME"

                # 获取该列的前3个样本值
                sample_values = df[col_name].dropna().head(3).tolist()
                if sample_values:
                    # 根据类型格式化样本值
                    if inferred_type == "TEXT":
                        sample_str = ", ".join(f"'{str(v)}'" for v in sample_values)
                    else:
                        sample_str = ", ".join(str(v) for v in sample_values)
                    sample_str = f"[{sample_str}]"
                else:
                    sample_str = "[]"

                columns_desc.append(f"{col_name} ({inferred_type}) {sample_str}")

            # 组合表信息
            info_text += f"表名: {table_name}\n"
            info_text += f"字段: {', '.join(columns_desc)}\n"
            info_text += "\n"

        return info_text.strip()

    def format_multitable_info(self) -> str:
        formatted = ""
        for table, columns in self.all_table_info.items():
            formatted += f"表名: {table}，字段: {', '.join(columns)}\n"
        return formatted

    def execute_sql(self, sql: str) -> Union[pd.DataFrame, str]:
        """
        执行SQL查询并返回结果

        Args:
            sql (str): 要执行的SQL查询语句。

        Returns:
            Union[pd.DataFrame, str]:
                - 如果查询成功，返回查询结果作为 pandas DataFrame。
                - 如果查询失败，返回错误信息的字符串。

        Raises:
            None: 错误通过返回值处理，而不是抛出异常。
            :param sql:
        """
        try:
            # 使用 pandas 执行查询并返回 DataFrame
            result_df = pd.read_sql_query(sql, self.connection)
            return result_df
        except Exception as e:
            # 返回格式化的错误信息
            return f"查询失败: {str(e)}"


if __name__ == "__main__":
    from nl2sql.model.embedding_model import EmbeddingModel

    apikey = "sk-68ac5f5ccf3540ba834deeeaecb48987"
    embd = EmbeddingModel(api_key=apikey, model='text-embedding-v2')

    data_path = '/Users/tiantiantian/Code/lab/engineering/nl2sql/examples/data/航空公司航班信息表.xlsx'
    excel_ds = ExcelDatasource(source_path=data_path, embedding_model=embd)

    result = excel_ds.find_closest_enum(input_value='华航')
    print(result)

    # excel_ds.get_column_enums(column_name='城市')

    # 假设用户问题提到了 “看下江城的效率”
    # print(excel_ds.find_closest_enum(input_value='江城'))

    # sql_result = excel_ds.execute_sql(sql='SELECT * FROM account_info_Sheet1')
    # print("执行SQL结果:", sql_result)
    #
    # data_info = excel_ds.get_data_info()
    # print(data_info)

